#ifndef _HYBTILEGRAPH_CUH__
#define _HYBTILEGRAPH_CUH__

#include "GraphData.cuh"
#include "config.hxx"

// #define DEBUG
#include "../utils.hxx"
#include "torch/extension.h"

namespace GraphData
{
  using namespace torch;

  template <typename T>
  struct chunk_t
  {
    // heavy chunk arts
    int *block_offset_sdat;
    int *block_offset_edat;
    int *block_offset_sb;
    int *h_tiled_edge;
    T *h_tiled_weight;
    int hcnt;
    int hdata_size;
    // light chunk arts
    int *block_offset_elst;
    short *l_src_data;
    int *l_dst_pos_sb;
    T *l_weight_data;
    int lcnt;
    int ldata_size;
    // dense chunk arts
    int *block_offset_sb_dense;
    T *d_mat_data;
    int dcnt;
    // stream buffer arts
    int *block_offset_dlst;
    short *agg_dst;
    int sbsize;

    int nv;
    int ne;
    bool weighted;
    int csize;
    int *reorder_map_src;
    int *reorder_map_dst;

    struct edge_t
    {
      int cid, v, u;
      T e;
    };

    chunk_t() : h_tiled_weight(nullptr), l_weight_data(nullptr), reorder_map_src(nullptr), reorder_map_dst(nullptr) {}

    void from_coo(int *src_list, int *dst_list, const T *weights,
                  const int nvertex, const int nedge, bool reord = false)
    {
      csize = CEIL(nvertex, BSZ);
      nv = nvertex;
      ne = nedge;
      weighted = (weights != nullptr);

      int nchunk = csize * csize;

      if (reord)
      {
        MYLOG("0. reodering")
        reorder_map_src = (int *)malloc(sizeof(int) * nvertex);
        reorder_map_dst = (int *)malloc(sizeof(int) * nvertex);
        reorder(src_list, dst_list, reorder_map_src, nvertex, nedge, true);
        reorder(src_list, dst_list, reorder_map_dst, nvertex, nedge, false);
      }
      else
      {
        reorder_map_src = reorder_map_dst = nullptr;
      }

      MYLOG("1. building chunks\n")
      MYLOG(" -- param check: THD=%d, BSZ=%d, H_THRESH=%d\n", THD, BSZ, H_THRESH);

      vector<edge_t> edge_coll;
      vector<int> ne(nchunk, 0);
      for (int e = 0; e < nedge; ++e)
      {
        int v = src_list[e];
        int u = dst_list[e];
        T ev = (weighted) ? weights[e] : (T)0;
        int row_id = v / BSZ;
        int col_id = u / BSZ;
        int cid = row_id * csize + col_id;

        edge_coll.push_back(edge_t{cid, v, u, ev});
        ++ne[cid];
      }

      sort(edge_coll.begin(), edge_coll.end(), [](const edge_t &a, const edge_t &b) -> bool
           {
             if (a.cid == b.cid)
               return a.v < b.v;
             else
               return a.cid < b.cid;
           });

      MYLOG("2. classify heavy & light\n");

      vector<int> ne_scan(nchunk + 1, 0);
      vector<int> &padded_ne = ne;

      vector<int> light;
      vector<int> heavy;
      vector<int> dense;

      for (int c = 0; c < nchunk; ++c)
      {
        ne_scan[c + 1] = ne_scan[c] + padded_ne[c];
        if (padded_ne[c] > D_THRESH)
          dense.push_back(c);
        else if (padded_ne[c] > H_THRESH)
          heavy.push_back(c);
        else
          light.push_back(c);
        MYLOG("  : chunk %d size = %d\n", c, padded_ne[c]);
        padded_ne[c] = ALIGN(padded_ne[c], 8);
      }
      dcnt = dense.size();
      hcnt = heavy.size();
      lcnt = light.size();

      create_heavy_chunk(edge_coll, ne_scan, heavy);
      create_dense_chunk(edge_coll, ne_scan, dense);
      create_lite_chunk(edge_coll, ne_scan, padded_ne, light);

      MYLOG("5. build stream buffe\n");
      sbsize = 0;
      vector<int> sbmat(nchunk, 0);
      block_offset_dlst = (int *)malloc(sizeof(int) * (csize + 1));
      memset(block_offset_dlst, 0, sizeof(int) * (csize + 1));

      for (int i = 0; i < nchunk; ++i)
      {
        sbmat[i] = (ne_scan[i + 1] - ne_scan[i] > H_THRESH) ? BSZ : padded_ne[i];
        sbsize += sbmat[i];
      }

      for (int c = 0; c < nchunk; ++c)
      {
        int col_panel = c % csize;
        block_offset_dlst[col_panel + 1] += sbmat[c];
      }
      for (int c = 0; c < csize; ++c)
        block_offset_dlst[c + 1] += block_offset_dlst[c];

      if (block_offset_dlst[csize] != sbsize)
        MYLOG(" -- fill block_offset_dlst error!\n");

      agg_dst = (short *)malloc(sizeof(short) * sbsize);
      memset(agg_dst, -1, sizeof(short) * sbsize);
#define TRANS(i, sz) ((i % sz) * sz + (i / sz))
      vector<int> sbcol_scan(nchunk + 1, 0);
      for (int c = 0; c < nchunk; ++c)
        sbcol_scan[c + 1] = sbmat[TRANS(c, csize)];
      for (int c = 0; c < nchunk; ++c)
      {
        sbcol_scan[c + 1] += sbcol_scan[c];
        MYLOG(" -- col_scan [%d] = %d\n", c + 1, sbcol_scan[c + 1]);
      }

      for (int i = 0; i < nchunk; ++i)
      {
        int base = sbcol_scan[TRANS(i, csize)];
        if (ne_scan[i + 1] - ne_scan[i] > H_THRESH)
        {
          for (int j = 0; j < BSZ; ++j)
            agg_dst[base + j] = j;
        }
        else
        {
          int elOffset = ne_scan[i];
          for (int j = 0; j < ne_scan[i + 1] - elOffset; ++j)
          {
            int u = edge_coll[elOffset + j].u;
            agg_dst[base + j] = (short)(u % BSZ);
          }
        }
      }

      int lbase = 0;
      for (uint64_t i = 0; i < light.size(); ++i)
      {
        int cid = light[i];
        int ne = padded_ne[cid];
        int cbase = sbcol_scan[TRANS(cid, csize)];
        for (int e = 0; e < ne; ++e)
        {
          if ((e & 7) == 0)
            l_dst_pos_sb[(lbase + e) >> 3] = cbase + e;
        }
        lbase += ne;
      }

      for (uint64_t i = 0; i < heavy.size(); ++i)
      {
        int cid = heavy[i];
        int cbase = sbcol_scan[TRANS(cid, csize)];

        block_offset_sb[i] = cbase;
      }

      for (uint64_t i = 0; i < dense.size(); ++i)
      {
        int cid = dense[i];
        int cbase = sbcol_scan[TRANS(cid, csize)];

        block_offset_sb_dense[i] = cbase;
      }
#undef TRANS
    }
#undef OMP_THD

    inline void to_tensor(Tensor &bosTn, Tensor &boeTn, Tensor &bosbTn, Tensor &hteTn, Tensor &htwTn,
                          Tensor &bosbdTn, Tensor &dmdTn,
                          Tensor &boelTn, Tensor &lslTn, Tensor &ldpsTn, Tensor &lwlTn,
                          Tensor &bodlTn, Tensor &adTn, Tensor &srcRordTn, Tensor &dstRordTn, TensorOptions &valOpt)
    {
      auto indOpt = TensorOptions().dtype(kInt32);
      auto sidOpt = TensorOptions().dtype(kInt16);

#define PACK(tn, ptr, size, ty, opt)                   \
  tn = torch::zeros((size), opt);                      \
  memcpy(tn.data_ptr<ty>(), ptr, (size) * sizeof(ty)); \
  free(ptr)

      if (this->reorder_map_dst)
      {
        PACK(srcRordTn, reorder_map_src, nv, int, indOpt);
        PACK(dstRordTn, reorder_map_dst, nv, int, indOpt);
      }

      if (dcnt > 0)
      {
        PACK(bosbdTn, block_offset_sb_dense, dcnt, int, indOpt);
        PACK(dmdTn, d_mat_data, BSZ * BSZ * dcnt, T, valOpt);
      }

      if (hcnt > 0)
      {
        PACK(bosTn, block_offset_sdat, hcnt, int, indOpt);
        PACK(boeTn, block_offset_edat, hcnt + 1, int, indOpt);
        PACK(bosbTn, block_offset_sb, hcnt, int, indOpt);
        PACK(hteTn, h_tiled_edge, hdata_size, int, indOpt);
      }

      if (lcnt > 0)
      {
        PACK(boelTn, block_offset_elst, csize + 1, int, indOpt);
        PACK(lslTn, l_src_data, ldata_size, short, sidOpt);
        PACK(ldpsTn, l_dst_pos_sb, CEIL(ldata_size, 8), int, indOpt);
      }

      PACK(bodlTn, block_offset_dlst, csize + 1, int, indOpt);
      PACK(adTn, agg_dst, sbsize, short, sidOpt);

      if (weighted)
      {
        if (hcnt > 0)
        {
          PACK(htwTn, h_tiled_weight, hdata_size, T, valOpt);
        }
        if (lcnt > 0)
        {
          PACK(lwlTn, l_weight_data, ldata_size, T, valOpt);
        }
      }
  #undef PACK
    }
          
    void cuda()
    {
      int *bosbd, *bos, *boe, *bosb, *hte, *boel, *ldps, *bodl;
      short *lsd, *ad;
      T *htw, *lwd, *dmd;

      H_ERR(cudaMalloc(&bosbd, sizeof(int)*dcnt));
      H_ERR(cudaMalloc(&bos, sizeof(int)*hcnt));
      H_ERR(cudaMalloc(&boe, sizeof(int)*(hcnt+1)));
      H_ERR(cudaMalloc(&bosb, sizeof(int)*hcnt));
      H_ERR(cudaMalloc(&hte, sizeof(int)*hdata_size));
      H_ERR(cudaMalloc(&boel, sizeof(int)*(csize+1)));
      H_ERR(cudaMalloc(&ldps, sizeof(int)*CEIL(ldata_size, 8)));
      H_ERR(cudaMalloc(&bodl, sizeof(int)*(csize+1)));

      H_ERR(cudaMalloc(&lsd, sizeof(short)*ldata_size));
      H_ERR(cudaMalloc(&ad, sizeof(short)*sbsize));

      if ( weighted ) {
        H_ERR(cudaMalloc(&htw, sizeof(T)*hdata_size));
        H_ERR(cudaMalloc(&lwd, sizeof(T)*ldata_size));
        H_ERR(cudaMalloc(&dmd, sizeof(T)*dcnt*BSZ*BSZ));
      }

      H_ERR(TODEV(bosbd, block_offset_sb_dense, dcnt));
      H_ERR(TODEV(bos, block_offset_sdat, hcnt));
      H_ERR(TODEV(boe, block_offset_edat, hcnt+1));
      H_ERR(TODEV(bosb, block_offset_sb, hcnt));
      H_ERR(TODEV(hte, h_tiled_edge, hdata_size));
      H_ERR(TODEV(boel, block_offset_elst, csize+1));
      H_ERR(TODEV(ldps, l_dst_pos_sb, CEIL(ldata_size, 8)));
      H_ERR(TODEV(bodl, block_offset_dlst, csize+1));
      H_ERR(TODEV(lsd, l_src_data, ldata_size));
      H_ERR(TODEV(ad, agg_dst, sbsize));

      if ( weighted ) {
        H_ERR(TODEV(htw, h_tiled_weight, hdata_size));
        H_ERR(TODEV(lwd, l_weight_data, ldata_size));
        H_ERR(TODEV(dmd, d_mat_data, dcnt*BSZ*BSZ));
      }

      free(block_offset_sb_dense);
      free(block_offset_sdat);
      free(block_offset_edat);
      free(block_offset_sb);
      free(block_offset_elst);
      free(block_offset_dlst);
      free(h_tiled_edge);
      free(l_src_data);
      free(l_dst_pos_sb);
      free(agg_dst);
      if ( weighted ) {
        free(h_tiled_weight);
        free(l_weight_data);
        free(d_mat_data);
      }

      block_offset_sb_dense = bosbd;
      block_offset_sdat = bos;
      block_offset_edat = boe;
      block_offset_sb = bosb;
      block_offset_elst = boel;
      block_offset_dlst = boel;
      h_tiled_edge = hte;
      l_src_data = lsd;
      l_dst_pos_sb = ldps;
      agg_dst = ad;
      if ( weighted ) {
        h_tiled_weight = htw;
        l_weight_data = lwd;
        d_mat_data = dmd;
      }
    }

    void cpu()
    {
      int *bosbd, *bos, *boe, *bosb, *hte, *boel, *ldps, *bodl;
      short *lsd, *ad;
      T *htw, *lwd, *dmd;

      bosbd = (int*) malloc(sizeof(int)*dcnt);
      bos = (int*) malloc(sizeof(int)*hcnt);
      boe = (int*) malloc(sizeof(int)*(hcnt+1));
      bosb = (int*) malloc(sizeof(int)*hcnt);
      hte = (int*) malloc(sizeof(int)*hdata_size);
      boel = (int*) malloc(sizeof(int)*(csize+1));
      ldps = (int*) malloc(sizeof(int)*CEIL(ldata_size, 8));
      bodl = (int*) malloc(sizeof(int)*(csize+1));

      lsd = (short*) malloc(sizeof(short)*ldata_size);
      ad = (short*) malloc(sizeof(short)*sbsize);

      if ( weighted ) {
        htw = (T *)malloc(sizeof(T) * hdata_size);
        lwd = (T *)malloc(sizeof(T) * ldata_size);
        dmd = (T *)malloc(sizeof(T) * BSZ * BSZ * dcnt);
      }

      H_ERR(TOHOST(block_offset_sb_dense, bosbd, dcnt));
      H_ERR(TOHOST(block_offset_sdat, bos, hcnt));
      H_ERR(TOHOST(block_offset_edat, boe, hcnt+1));
      H_ERR(TOHOST(block_offset_sb, bosb, hcnt));
      H_ERR(TOHOST(h_tiled_edge, hte, hdata_size));
      H_ERR(TOHOST(block_offset_elst, boel, csize+1));
      H_ERR(TOHOST(l_dst_pos_sb, ldps, CEIL(ldata_size, 8)));
      H_ERR(TOHOST(block_offset_dlst, bodl, csize+1));
      H_ERR(TOHOST(l_src_data, lsd, ldata_size));
      H_ERR(TOHOST(agg_dst, ad, sbsize));
      if ( weighted ) {
        H_ERR(TOHOST(l_weight_data, lwd, ldata_size));
        H_ERR(TOHOST(h_tiled_weight, htw, hdata_size));
        H_ERR(TOHOST(d_mat_data, dmd, dcnt * BSZ * BSZ));
      }

      cudaFree(block_offset_sb_dense);
      cudaFree(block_offset_sdat);
      cudaFree(block_offset_edat);
      cudaFree(block_offset_sb);
      cudaFree(block_offset_elst);
      cudaFree(block_offset_dlst);
      cudaFree(h_tiled_edge);
      cudaFree(l_src_data);
      cudaFree(l_dst_pos_sb);
      cudaFree(agg_dst);
      if ( weighted ) {
        cudaFree(l_weight_data);
        cudaFree(h_tiled_edge);
        cudaFree(d_mat_data);
      }

      block_offset_sb_dense = bosbd;
      block_offset_sdat = bos;
      block_offset_edat = boe;
      block_offset_sb = bosb;
      block_offset_elst = boel;
      block_offset_dlst = boel;
      h_tiled_edge = hte;
      l_src_data = lsd;
      l_dst_pos_sb = ldps;
      agg_dst = ad;
      if ( weighted ) {
        l_weight_data = lwd;
        h_tiled_weight = htw;
        d_mat_data = dmd;
      }
    }

  private:
    void create_dense_chunk(vector<edge_t> &edge_coll, vector<int> &ne_scan, vector<int> &dense)
    {
      MYLOG("3.5. build dense chunks\n")
      block_offset_sb_dense = (int *)malloc(sizeof(int) * dcnt);
      d_mat_data = (T *)malloc(sizeof(T) * BSZ * BSZ * dcnt);

      memset(d_mat_data, 0, sizeof(T) * BSZ * BSZ * dcnt);

      for (uint64_t c = 0; c < dense.size(); ++c)
      {
        int cid = dense[c];
        int dm_base = BSZ * BSZ * c;

        for (int e = ne_scan[cid]; e < ne_scan[cid + 1]; ++e)
        {
          int col = edge_coll[e].v % BSZ;
          int row = edge_coll[e].u % BSZ;
          // saving the mat in a transposed style
          d_mat_data[dm_base + col * BSZ + row] = edge_coll[e].e;
        }

        MYLOG(" -- dense tile %d (cid=%d) with offset %d\n", c, cid, dm_base);
      }
    }

    void create_lite_chunk(vector<edge_t> &edge_coll, vector<int> &ne_scan, vector<int> &padded_ne, vector<int> &light)
    {
      MYLOG("4. build light chunks\n")
      ldata_size = 0;
      if (lcnt > 0)
      {
        block_offset_elst = (int *)malloc(sizeof(int) * (csize + 1));
        memset(block_offset_elst, 0, sizeof(int) * (csize + 1));

        for (int i = 0; i < lcnt; ++i)
        {
          int cid = light[i];
          int row_panel = cid / csize;

          block_offset_elst[row_panel + 1] += padded_ne[cid];
          MYLOG("  : lite tile %d (%d, %d) -- size=%d\n", cid, row_panel, cid % csize, padded_ne[cid]);
        }
        for (int i = 0; i < csize; ++i)
          block_offset_elst[i + 1] += block_offset_elst[i];

        ldata_size = block_offset_elst[csize];
        MYLOG(" -- ldata_size = %d\n", ldata_size);

        l_src_data = (short *)malloc(sizeof(short) * ldata_size);
        l_dst_pos_sb = (int *)malloc(sizeof(int) * CEIL(ldata_size, 8));
        l_weight_data = (weighted) ? (T *)malloc(sizeof(T) * ldata_size) : nullptr;

        memset(l_src_data, -1, sizeof(short) * ldata_size);
        int top = 0;
        for (uint64_t i = 0; i < light.size(); ++i)
        {
          int cid = light[i];
          for (int e = ne_scan[cid]; e < ne_scan[cid + 1]; ++e)
          {
            int v = edge_coll[e].v;
            T ew = edge_coll[e].e;

            l_src_data[top] = (short)(v % BSZ);
            if (weighted)
              l_weight_data[top] = ew;
            ++top;
          }
          MYLOG("  : lite tile %d (id=%d) -- cur top=%d ", i, cid, top);
          top += padded_ne[cid] - ne_scan[cid + 1] + ne_scan[cid];
          MYLOG("padded to %d (%d ->%d)\n", top, ne_scan[cid + 1] - ne_scan[cid], padded_ne[cid]);
        }
        if (top != ldata_size)
        {
          MYLOG(" -- fill l_src_data error: expect=%d got=%d\n", ldata_size, top);
          exit(1);
        }
      }
    }

    void create_heavy_chunk(vector<edge_t> &edge_coll, vector<int> &ne_scan, vector<int> &heavy)
    {
      MYLOG(" -- chunck type stat: heavy=%d, lite=%d\n", hcnt, lcnt);

      MYLOG("3. build heavy chunks\n");
      hdata_size = 0;
      if (hcnt > 0)
      {
        int top = 0;
        block_offset_sdat = (int *)malloc(hcnt * sizeof(int));

        for (int h = 0; h < hcnt; ++h)
          block_offset_sdat[h] = heavy[h] / csize;
        vector<int> block_size_edat(hcnt, 0);

        block_offset_edat = (int *)malloc(sizeof(int) * (hcnt + 1));
        block_offset_sb = (int *)malloc(sizeof(int) * (hcnt));
        for (int i = 0; i < hcnt; ++i)
        {
          int idx = heavy[i];
          MYLOG("  : heavy tile %d(%d~%d in %d)\n", top, ne_scan[idx], ne_scan[idx + 1], edge_coll.size());
          sort(edge_coll.begin() + ne_scan[idx], edge_coll.begin() + ne_scan[idx + 1],
               [](const edge_t &a, const edge_t &b) -> bool
               {
                 if (a.u == b.u)
                   return a.v < b.v;
                 else
                   return a.u < b.u;
               });

          block_size_edat[top] = 1 + CEIL(ne_scan[idx + 1] - ne_scan[idx] + BSZ, THD);
          block_offset_edat[top] = hdata_size;
          hdata_size += block_size_edat[top++] * THD;
          MYLOG(" -- size=%d*THD offset=%d; cur hsize=%d\n", block_size_edat[top - 1], block_offset_edat[top - 1], hdata_size);
        }

        block_offset_edat[top] = hdata_size;
        h_tiled_edge = (int *)malloc(sizeof(int) * hdata_size);
        h_tiled_weight = (weighted) ? (T *)malloc(sizeof(T) * hdata_size) : nullptr;

        memset(h_tiled_edge, -1, sizeof(int) * hdata_size);

#define HPOS(i, dim) ((i % dim) * THD + (i / dim))
        for (int i = 0; i < hcnt; ++i)
        {
          int cid = heavy[i];
          int top = 0;
          int cur = -1;
          int base = block_offset_edat[i];
          int rows = block_size_edat[i];

          for (int e = ne_scan[cid]; e < ne_scan[cid + 1]; ++e)
          {
            int v = edge_coll[e].v;
            int u = edge_coll[e].u;
            T ew = edge_coll[e].e;

            if (cur != u)
            {
              h_tiled_edge[base + HPOS(top, rows)] = HASH(u % BSZ);
              cur = u;
              ++top;
            }
            if (top % rows == 0)
            {
              h_tiled_edge[base + HPOS(top, rows)] = HASH(cur % BSZ);
              ++top;
            }

            h_tiled_edge[base + HPOS(top, rows)] = v % BSZ;
            if (weighted)
              h_tiled_weight[base + HPOS(top, rows)] = ew;
            ++top;
          }

          if (top > rows * THD)
          {
            MYLOG(" -- fill error, exit %d %d %d\n", top, rows * THD, (ne_scan[cid + 1] - ne_scan[cid]));
            exit(1);
          }
        }
#undef HPOS
        // profile step, optional
        int valid_edge = 0;
        for (int i = 0; i < hdata_size; ++i)
          if (h_tiled_edge[i] > 0)
            ++valid_edge;
        MYLOG(" -- valid vertices: %d\n", valid_edge);
      }
    }
  };
};

#endif